package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	crand "crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"net/url"
	"strings"
)

func RandomString(length int) string {
	return string(RandomBytes(length))
}

func RandomBytes(length int) []byte {
	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	b := make([]byte, length)
	for i := range b {
		b[i] = charset[rand.Intn(len(charset))]
	}
	return b
}

func HttpGet(full_url string, params url.Values) ([]byte, error) {
	url, err := url.Parse(full_url)
	if err != nil {
		return nil, err
	}
	url.RawQuery = params.Encode()
	resp, err := http.Get(url.String())
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return body, nil
}

func HttpPost(full_url string, req_str string, key string, iv string, token string) ([]byte, error) {
	request, err := http.NewRequest("POST", full_url, strings.NewReader(req_str))
	if err != nil {
		fmt.Println("http new request error")
		return nil, err
	}
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Connection", "Keep-Alive")
	request.Header.Set("CRYPTO", key)
	request.Header.Set("IV", iv)
	request.Header.Set("SIGN", "sign")
	if len(token) > 0 {
		request.Header.Set("token", "bearer "+token)
	}
	client := &http.Client{}
	var resp *http.Response
	resp, err = client.Do(request)
	if err != nil {
		fmt.Println("http do error")
		return nil, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		fmt.Println("read resp error")
		return nil, err
	}
	return body, nil
}

func AesEncrypt(plaintext []byte, key []byte) (crypted []byte, iv []byte, err error) {
	block, err := aes.NewCipher(key) // 创建新的AES cipher对象
	if err != nil {
		return nil, nil, err
	}

	blockSize := block.BlockSize()
	paddingBytes := pkcs7Padding(plaintext, blockSize) // 填充

	crypted = make([]byte, len(paddingBytes))
	iv = RandomBytes(blockSize) // 创建长度与blockSize一致的iv

	cbcMode := cipher.NewCBCEncrypter(block, iv) // 创建CBC模式的处理器
	cbcMode.CryptBlocks(crypted, paddingBytes)   // 执行加密

	return crypted, iv, nil
}

func AesDecrypt(ciphertext []byte, key []byte, iv []byte) (crypted []byte, err error) {
	// 创建实例
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	// 获取块的大小
	// blockSize := block.BlockSize()
	// 使用cbc
	blockMode := cipher.NewCBCDecrypter(block, iv)
	// 初始化解密数据接收切片
	crypted = make([]byte, len(ciphertext))
	// 执行解密
	blockMode.CryptBlocks(crypted, ciphertext)
	// 去除填充
	crypted = pkcs7UnPadding(crypted)
	return crypted, nil
}

// pkcs7Padding 填充
func pkcs7Padding(data []byte, blockSize int) []byte {
	// 判断缺少几位长度. 最少1, 最多blockSize
	padding := blockSize - len(data)%blockSize
	// 补足位数. 把切片[]byte{byte(padding)}复制padding个
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// pkcs7UnPadding 填充的反向操作
func pkcs7UnPadding(data []byte) []byte {
	length := len(data)
	unpaddingLength := int(data[length-1])
	return data[:(length - unpaddingLength)]
}

func generate_key(pubKey *rsa.PublicKey) (key []byte, base64Key string) {
	key = RandomBytes(32)
	cipherKey, err := rsa.EncryptPKCS1v15(crand.Reader, pubKey, key)
	if err != nil {
		fmt.Println("Failed to encrypt data with RSA")
		return
	}
	base64Key = base64.StdEncoding.EncodeToString(cipherKey)
	return key, base64Key
}

func Base64_iv(pub *rsa.PublicKey, iv []byte) string {
	cipherIv, err := rsa.EncryptPKCS1v15(crand.Reader, pub, iv)
	if err != nil {
		fmt.Println("Failed to encrypt data with RSA", err)
		return ""
	}
	base64Iv := base64.StdEncoding.EncodeToString(cipherIv)
	return base64Iv
}

func Encrypt(request_body any, key []byte) (string, []byte, error) {
	json_bytes, err := json.Marshal(request_body)
	if err != nil {
		fmt.Println("json marshal error")
		return "", nil, err
	}
	cipher, iv, err := AesEncrypt(json_bytes, key)
	if err != nil {
		fmt.Println("aes encrypt error")
		return "", nil, err
	}
	base64 := base64.StdEncoding.EncodeToString(cipher)
	return base64, iv, nil
}
